import os
import datetime
import tables
import numpy as np
from matplotlib.figure import Figure
from matplotlib.dates import DateFormatter
from matplotlib.ticker import MultipleLocator
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.mlab import movavg
from cwd.lib.data import PrecipGrid
from cwd.lib.utils import TextOutput

import pdb
PDB = pdb.set_trace

class PrecipBarGraph(object):
    def __init__(self, lat, lon, month, span=1, level=0, runavg=0, data=None):
        self.lat = lat
        self.lon = lon
        self.level = int(level)
        self.month = month
        self.span = span
        self.runavg = runavg

        if data is None:
            grid = PrecipGrid(lat,lon,level)
            self.last_day = grid.last_day
            self.level_map = grid.level_map
            self.month_data = grid.get_monthly_data(month, span)

    def write_to(self,output,format="png"):
        """
        Writes image to file-like 'output'
        """
        if format == "text":
            self.graph_ax(output, format, data=self.month_data, month=self.month, span=self.span)
            output.get_text()
        else:
            # Create figure (the plot container)
            fig = Figure(figsize=(9,6), facecolor='w')
            # Pass all parameters to graphing function (alters 'fig')
            self.graph_ax(fig, format, data=self.month_data, month=self.month, span=self.span)
            # Create a canvas for png write
            canvas = FigureCanvas(fig)
            canvas.print_png(output)

    def graph_ax(self, outobj, format, *args, **kwargs):
        """
        Add a subplot to outobj of freezing levels for a month.
        """
        if not kwargs:
                return
        month = self.month
        month_data = self.month_data
        span = self.span
        runavg = self.runavg
        if runavg % 2: #odd 
            s_i = int(np.ceil(runavg/2.0)) - 1
            e_i = s_i
        else: #even 
            s_i = (runavg/2) - 1
            e_i = s_i + 1
        level = self.level
        years = np.arange(1948,1948+month_data.size)
        if month > datetime.datetime.now().month: # Don't include months that haven't yet passed
            years = years[:-1]
            month_data = month_data[:-1]

        # Find the mean from 1970-2000 = index 22:53
        # Find the mean from 1981-2010 = index 33:63
        mean = month_data[33:63].mean()
        inv_data = mean - month_data

        if format != "text":
            ax = outobj.add_axes([0.09, 0.15, .90, 0.78])

            # Insert titles
            ax.set_ylabel("% Precip as Snow")
            if ( month - span < 0 ): xlabel = "Ending Year"
            else: xlabel = "Year"
            ax.set_xlabel("%s (%d - %d)" % (xlabel, years[0], years[-1]))
            outobj.text(0.745, 0.047, "Western Regional Climate Center",
                        {'size':"x-small", 'family':"monospace"})
            outobj.text(0.830, 0.070, "Means from 1981-2010",
                        {'size':"x-small", 'family':"monospace"})
            outobj.text(0.79, 0.090, "Last data: %s" % self.last_day.strftime("%Y-%m-%d"),
                        {'size':"small", 'family':"monospace"})
            if span > 1:
                title = "%d Months Ending in %s" % (span, datetime.date(2000,month,1).strftime("%B"))
            else:
                title = "%s" % (datetime.date(2000,month,1).strftime("%B"))
            title += " % of Precip as Snow"
            title += u" %4.2f\u00b0N, %4.2f\u00b0W %dm" % (self.lat, self.lon, self.level_map[level])
            ax.set_title(title, family='serif')

            # Manually scale the axes
            ax.set_autoscale_on(False)
            ax.set_xbound((years[0]-1, years[-1]+1))
            y_min, y_max = max(np.nanmin(month_data)-5, 0), min(np.nanmax(month_data)+5, 100)
            ax.set_ybound(y_min, y_max)
            ax.yaxis.set_major_locator(MultipleLocator(10))
            ax.yaxis.set_minor_locator(MultipleLocator(5))
            ax.xaxis.set_major_locator(MultipleLocator(10))
            ax.xaxis.set_minor_locator(MultipleLocator(1))

            # graph the mean line
            ax.axhline(y=mean,color="black")
            # graph one std deviation
            ax.axhline(y=mean - month_data.std(),color="#999999", linestyle=":")
            ax.axhline(y=mean + month_data.std(),color="#999999", linestyle=":")
            # graph the bars below the mean
            ax.bar(years[month_data < mean], inv_data[month_data < mean], color="red", align="center", bottom=month_data[month_data<mean])
            # graph the bars above the mean
            inv_data = month_data - mean
            ax.bar(years[month_data >= mean], inv_data[month_data >= mean], color="blue", align="center", bottom=mean)
            if runavg:
                if e_i == 0:
                    x_years = years
                else:
                    x_years = years[s_i:-e_i]

                runavg_line = ax.plot(x_years, movavg(month_data,runavg), color='black', linewidth=2, label="%d year average"% runavg)
                ax.legend(loc=(0, -0.125))
        else:
            outobj.add(years, header="year")
            outobj.add(month_data, header="percent")
            outobj.add([mean]*years.size, header="mean")
            if runavg:
                d = [-9999]*(s_i) + list(movavg(month_data,float(runavg))) + [-9999]*e_i
                outobj.add( d, header="movavg_%d"%runavg )

if __name__ == '__main__':
    lat = 37.5
    lon = 120
    level = 5
    month = 8
    span = 1
    runavg = 0
    text_output = TextOutput(delim="  ")
    bargraph = PrecipBarGraph(lat,lon,month,span=span,level=level,runavg=runavg)
    bargraph.write_to(text_output,format="text")
